# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
# ============================================================================
# Copyright 2021 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""Layer Normalization layer."""

import tensorflow.compat.v2 as tf
# pylint: disable=g-classes-have-attributes

from keras import constraints
from keras import initializers
from keras import regularizers
from keras.engine.base_layer import Layer
# from keras.utils import tf_utils
import npu_device

# from tensorflow.python.util.tf_export import keras_export

def validate_axis(axis, input_shape):
  """Validate an axis value and returns its standardized form.
  Args:
    axis: Value to validate. Can be an integer or a list/tuple of integers.
      Integers may be negative.
    input_shape: Reference input shape that the axis/axes refer to.
  Returns:
    Normalized form of `axis`, i.e. a list with all-positive values.
  """
  input_shape = tf.TensorShape(input_shape)
  rank = input_shape.rank
  if not rank:
    raise ValueError(
        f'Input has undefined rank. Received: input_shape={input_shape}')

  # Convert axis to list and resolve negatives
  if isinstance(axis, int):
    axis = [axis]
  else:
    axis = list(axis)
  for idx, x in enumerate(axis):
    if x < 0:
      axis[idx] = rank + x

  # Validate axes
  for x in axis:
    if x < 0 or x >= rank:
      raise ValueError(
          'Invalid value for `axis` argument. '
          'Expected 0 <= axis < inputs.rank (with '
          f'inputs.rank={rank}). Received: axis={tuple(axis)}')
  if len(axis) != len(set(axis)):
    raise ValueError(f'Duplicate axis: {tuple(axis)}')
  return axis


class LayerNormalization(Layer):
    """Layer normalization layer (Ba et al., 2016).
    Normalize the activations of the previous layer for each given example in a
    batch independently, rather than across a batch like Batch Normalization.
    i.e. applies a transformation that maintains the mean activation within each
    example close to 0 and the activation standard deviation close to 1.
    Given a tensor `inputs`, moments are calculated and normalization
    is performed across the axes specified in `axis`.
    Example:
    >>> data = tf.constant(np.arange(10).reshape(5, 2) * 10, dtype=tf.float32)
    >>> print(data)
    tf.Tensor(
    [[ 0. 10.]
     [20. 30.]
     [40. 50.]
     [60. 70.]
     [80. 90.]], shape=(5, 2), dtype=float32)
    >>> layer = tf.keras.layers.LayerNormalization(axis=1)
    >>> output = layer(data)
    >>> print(output)
    tf.Tensor(
    [[-1. 1.]
     [-1. 1.]
     [-1. 1.]
     [-1. 1.]
     [-1. 1.]], shape=(5, 2), dtype=float32)
    Notice that with Layer Normalization the normalization happens across the
    axes *within* each example, rather than across different examples in the
    batch.
    If `scale` or `center` are enabled, the layer will scale the normalized
    outputs by broadcasting them with a trainable variable `gamma`, and center
    the outputs by broadcasting with a trainable variable `beta`. `gamma` will
    default to a ones tensor and `beta` will default to a zeros tensor, so that
    centering and scaling are no-ops before training has begun.
    So, with scaling and centering enabled the normalization equations
    are as follows:
    Let the intermediate activations for a mini-batch to be the `inputs`.
    For each sample `x_i` in `inputs` with `k` features, we compute the mean and
    variance of the sample:
    ```python
    mean_i = sum(x_i[j] for j in range(k)) / k
    var_i = sum((x_i[j] - mean_i) ** 2 for j in range(k)) / k
    ```
    and then compute a normalized `x_i_normalized`, including a small factor
    `epsilon` for numerical stability.
    ```python
    x_i_normalized = (x_i - mean_i) / sqrt(var_i + epsilon)
    ```
    And finally `x_i_normalized ` is linearly transformed by `gamma` and `beta`,
    which are learned parameters:
    ```python
    output_i = x_i_normalized * gamma + beta
    ```
    `gamma` and `beta` will span the axes of `inputs` specified in `axis`, and
    this part of the inputs' shape must be fully defined.
    For example:
    >>> layer = tf.keras.layers.LayerNormalization(axis=[1, 2, 3])
    >>> layer.build([5, 20, 30, 40])
    >>> print(layer.beta.shape)
    (20, 30, 40)
    >>> print(layer.gamma.shape)
    (20, 30, 40)
    Note that other implementations of layer normalization may choose to define
    `gamma` and `beta` over a separate set of axes from the axes being
    normalized across. For example, Group Normalization
    ([Wu et al. 2018](https://arxiv.org/abs/1803.08494)) with group size of 1
    corresponds to a Layer Normalization that normalizes across height, width,
    and channel and has `gamma` and `beta` span only the channel dimension.
    So, this Layer Normalization implementation will not match a Group
    Normalization layer with group size set to 1.
    Args:
      axis: Integer or List/Tuple. The axis or axes to normalize across. Typically
        this is the features axis/axes. The left-out axes are typically the batch
        axis/axes. This argument defaults to `-1`, the last dimension in the
        input.
      epsilon: Small float added to variance to avoid dividing by zero. Defaults
        to 1e-3
      center: If True, add offset of `beta` to normalized tensor. If False, `beta`
        is ignored. Defaults to True.
      scale: If True, multiply by `gamma`. If False, `gamma` is not used. Defaults
        to True. When the next layer is linear (also e.g. `nn.relu`), this can be
        disabled since the scaling will be done by the next layer.
      beta_initializer: Initializer for the beta weight. Defaults to zeros.
      gamma_initializer: Initializer for the gamma weight. Defaults to ones.
      beta_regularizer: Optional regularizer for the beta weight. None by default.
      gamma_regularizer: Optional regularizer for the gamma weight. None by
        default.
      beta_constraint: Optional constraint for the beta weight. None by default.
      gamma_constraint: Optional constraint for the gamma weight. None by default.
    Input shape:
      Arbitrary. Use the keyword argument `input_shape` (tuple of
      integers, does not include the samples axis) when using this layer as the
      first layer in a model.
    Output shape:
      Same shape as input.
    Reference:
      - [Lei Ba et al., 2016](https://arxiv.org/abs/1607.06450).
    """

    def __init__(self,
                 axis=-1,
                 epsilon=1e-3,
                 center=True,
                 scale=True,
                 beta_initializer='zeros',
                 gamma_initializer='ones',
                 beta_regularizer=None,
                 gamma_regularizer=None,
                 beta_constraint=None,
                 gamma_constraint=None,
                 **kwargs):
        super(LayerNormalization, self).__init__(**kwargs)
        if isinstance(axis, (list, tuple)):
            self.axis = list(axis)
        elif isinstance(axis, int):
            self.axis = axis
        else:
            raise TypeError('Expected an int or a list/tuple of ints for the '
                            'argument \'axis\', but received: %r' % axis)

        self.epsilon = epsilon
        self.center = center
        self.scale = scale
        self.beta_initializer = initializers.get(beta_initializer)
        self.gamma_initializer = initializers.get(gamma_initializer)
        self.beta_regularizer = regularizers.get(beta_regularizer)
        self.gamma_regularizer = regularizers.get(gamma_regularizer)
        self.beta_constraint = constraints.get(beta_constraint)
        self.gamma_constraint = constraints.get(gamma_constraint)

        self.supports_masking = True

    def build(self, input_shape):
        self.axis = validate_axis(self.axis, input_shape)
        input_shape = tf.TensorShape(input_shape)
        rank = input_shape.rank

        param_shape = [input_shape[dim] for dim in self.axis]
        if self.scale:
            self.gamma = self.add_weight(
                name='gamma',
                shape=param_shape,
                initializer=self.gamma_initializer,
                regularizer=self.gamma_regularizer,
                constraint=self.gamma_constraint,
                trainable=True,
                experimental_autocast=False)
        else:
            self.gamma = None

        if self.center:
            self.beta = self.add_weight(
                name='beta',
                shape=param_shape,
                initializer=self.beta_initializer,
                regularizer=self.beta_regularizer,
                constraint=self.beta_constraint,
                trainable=True,
                experimental_autocast=False)
        else:
            self.beta = None

        self.built = True

    def call(self, inputs):
        inputs = tf.cast(inputs, self.compute_dtype)
        # begin_params_axis 从第几轴开始,inputs和gama/beta相同
        # begin_norm_axis对第几轴开始进行reduce,影响输出
        outputs, mean, variance = npu_device.gen_npu_ops.fused_layer_norm(inputs,
                                                                          self.gamma,
                                                                          self.beta,
                                                                          begin_norm_axis=self.axis[0],
                                                                          begin_params_axis=-1,
                                                                          epsilon=1e-12,
                                                                          name="fused_layer_norm")

        return outputs

    def compute_output_shape(self, input_shape):
        return input_shape

    def get_config(self):
        config = {
            'axis': self.axis,
            'epsilon': self.epsilon,
            'center': self.center,
            'scale': self.scale,
            'beta_initializer': initializers.serialize(self.beta_initializer),
            'gamma_initializer': initializers.serialize(self.gamma_initializer),
            'beta_regularizer': regularizers.serialize(self.beta_regularizer),
            'gamma_regularizer': regularizers.serialize(self.gamma_regularizer),
            'beta_constraint': constraints.serialize(self.beta_constraint),
            'gamma_constraint': constraints.serialize(self.gamma_constraint)
        }
        base_config = super(LayerNormalization, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))
